#!/bin/bash
#
# setup-dual-user.sh - 将单用户 Arch Linux 系统转换为双用户方案
#
# 功能：
# - 将 @home 子卷重命名为 @home_<主用户>
# - 创建 @home_test 和 @home_shared 子卷
# - 创建测试用户
# - 配置共享目录权限
# - 更新 /etc/fstab
# - (可选) 安装 niri
#
# 使用方法：
# 1. 交互式执行（推荐）：
#    sudo bash setup-dual-user.sh
#
# 2. 修改默认配置后执行：
#    编辑脚本，修改下方的 "配置变量" 区域
#    DEFAULT_TEST_USER="myniri"  # 修改默认测试用户名
#    然后执行：sudo bash setup-dual-user.sh
#
# 3. 查看详细说明：
#    cat setup-dual-user-使用说明.md
#
# 作者：Claude & Lin
# 日期：2025-11-05
# 版本：v1.1
#

set -e  # 遇到错误立即退出

##############################################################################
# 颜色和输出函数
##############################################################################

RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
BLUE='\033[0;34m'
CYAN='\033[0;36m'
NC='\033[0m' # No Color

log_info() {
    echo -e "${GREEN}[INFO]${NC} $1"
}

log_step() {
    echo -e "${CYAN}[步骤]${NC} $1"
}

log_warn() {
    echo -e "${YELLOW}[警告]${NC} $1"
}

log_error() {
    echo -e "${RED}[错误]${NC} $1"
}

log_success() {
    echo -e "${GREEN}[✓]${NC} $1"
}

##############################################################################
# 配置变量（可在脚本执行前修改）
##############################################################################
#
# 说明：
# - 这些是默认值，脚本运行时会询问用户确认或修改
# - 如果想避免交互，可以直接修改这些值
#
# 示例：
#   DEFAULT_TEST_USER="niri-test"   # 使用自定义的测试用户名
#   DEFAULT_INSTALL_NIRI="yes"      # 默认安装 niri
#
##############################################################################

# 测试用户名（默认值，脚本运行时可修改）
DEFAULT_TEST_USER="test"

# 是否安装 niri（默认 no，脚本运行时可修改）
DEFAULT_INSTALL_NIRI="no"

##############################################################################
# 全局变量（脚本运行时自动设置）
##############################################################################

MAIN_USER=""           # 主用户名（自动检测）
TEST_USER=""           # 测试用户名（从用户输入获取）
ROOT_DISK=""           # 根分区磁盘
BTRFS_ROOT=""          # Btrfs 根分区
SNAPSHOT_ID=""         # Snapper 快照 ID
TEMP_MNT=""            # 临时挂载点
INSTALL_NIRI="no"      # 是否安装 niri
LOG_FILE="/var/log/setup-dual-user.log"

##############################################################################
# 1. 前置检查
##############################################################################

check_root() {
    log_step "检查 root 权限..."
    if [ "$EUID" -ne 0 ]; then
        log_error "请使用 sudo 或 root 身份运行此脚本"
        exit 1
    fi
    log_success "root 权限检查通过"
}

check_btrfs() {
    log_step "检查文件系统类型..."
    
    if ! findmnt -n -o FSTYPE / | grep -q "btrfs"; then
        log_error "根文件系统不是 Btrfs"
        log_error "此脚本仅支持 Btrfs 文件系统"
        exit 1
    fi
    
    # 获取根分区设备
    ROOT_DISK=$(findmnt -n -o SOURCE /)
    BTRFS_ROOT=$(echo "$ROOT_DISK" | sed 's/\[.*\]//')  # 移除 [subvol] 部分
    
    log_success "文件系统: Btrfs"
    log_info "根分区: $BTRFS_ROOT"
}

check_already_dual_user() {
    log_step "检查是否已经是双用户结构..."
    
    # 检查是否存在 @home_shared 子卷
    TEMP_CHECK=$(mktemp -d)
    mount "$BTRFS_ROOT" "$TEMP_CHECK"
    
    if [ -d "$TEMP_CHECK/@home_shared" ]; then
        umount "$TEMP_CHECK"
        rmdir "$TEMP_CHECK"
        log_error "系统已经是双用户结构"
        log_error "检测到 @home_shared 子卷存在"
        log_info "如需重新配置，请先手动回滚"
        exit 1
    fi
    
    umount "$TEMP_CHECK"
    rmdir "$TEMP_CHECK"
    log_success "当前是单用户结构，可以转换"
}

detect_main_user() {
    log_step "检测主用户..."
    
    # 获取 /home 下的用户（排除 shared, lost+found 等）
    local users=()
    for dir in /home/*; do
        if [ -d "$dir" ]; then
            local username=$(basename "$dir")
            # 排除特殊目录
            if [[ "$username" != "lost+found" && "$username" != "shared" ]]; then
                users+=("$username")
            fi
        fi
    done
    
    if [ ${#users[@]} -eq 0 ]; then
        log_error "未检测到任何用户"
        log_error "请确保系统已经创建了主用户"
        exit 1
    elif [ ${#users[@]} -eq 1 ]; then
        MAIN_USER="${users[0]}"
        log_success "检测到主用户: $MAIN_USER"
    else
        log_warn "检测到多个用户: ${users[*]}"
        log_info "请选择主用户（将保留其数据）:"
        select user in "${users[@]}"; do
            if [ -n "$user" ]; then
                MAIN_USER="$user"
                break
            fi
        done
        log_info "已选择主用户: $MAIN_USER"
    fi
}

check_snapper() {
    log_step "检查 Snapper 是否可用..."
    
    if ! command -v snapper &> /dev/null; then
        log_warn "Snapper 未安装"
        log_warn "建议安装 Snapper 以便创建快照保护"
        read -p "是否继续（不创建快照）？ [y/N] " -n 1 -r
        echo
        if [[ ! $REPLY =~ ^[Yy]$ ]]; then
            log_info "已取消操作"
            exit 0
        fi
    else
        log_success "Snapper 可用"
    fi
}

check_system() {
    echo ""
    echo "========================================"
    echo "  双用户方案转换 - 系统检查"
    echo "========================================"
    echo ""
    
    check_root
    check_btrfs
    check_already_dual_user
    detect_main_user
    check_snapper
    
    echo ""
    log_success "所有检查通过！"
    echo ""
}

##############################################################################
# 2. 信息收集
##############################################################################

collect_info() {
    echo "========================================"
    echo "  配置信息收集"
    echo "========================================"
    echo ""
    
    # 主用户（已检测）
    log_info "主用户: $MAIN_USER"
    read -p "是否正确？ [Y/n] " -n 1 -r
    echo
    if [[ $REPLY =~ ^[Nn]$ ]]; then
        read -p "请输入主用户名: " MAIN_USER
    fi
    
    # 测试用户名
    echo ""
    read -p "测试用户名 [默认: $DEFAULT_TEST_USER]: " input_test_user
    TEST_USER=${input_test_user:-$DEFAULT_TEST_USER}
    
    # 检查用户名冲突
    if id "$TEST_USER" &>/dev/null; then
        log_error "用户 $TEST_USER 已存在"
        exit 1
    fi
    
    if [ "$TEST_USER" = "$MAIN_USER" ]; then
        log_error "测试用户名不能与主用户相同"
        exit 1
    fi
    
    log_info "测试用户: $TEST_USER"
    
    # 是否安装 niri
    echo ""
    read -p "是否安装 niri Wayland compositor？ [y/N] " -n 1 -r
    echo
    if [[ $REPLY =~ ^[Yy]$ ]]; then
        INSTALL_NIRI="yes"
        
        # 检查 yay
        if ! command -v yay &> /dev/null; then
            log_warn "未检测到 yay (AUR helper)"
            log_info "安装 niri 需要 yay，脚本将跳过 niri 安装"
            log_info "您可以稍后手动安装: yay -S niri"
            INSTALL_NIRI="no"
        fi
    fi
    
    # 确认信息
    echo ""
    echo "========================================"
    echo "  配置摘要"
    echo "========================================"
    echo "  主用户:       $MAIN_USER (保留数据)"
    echo "  测试用户:     $TEST_USER (新建)"
    echo "  安装 niri:    $INSTALL_NIRI"
    echo "  根分区:       $BTRFS_ROOT"
    echo ""
    echo "将执行的操作:"
    echo "  1. 创建 Snapper 快照（保护数据）"
    echo "  2. 重构子卷："
    echo "     @home           → @home_$MAIN_USER"
    echo "     @home_test      (新建)"
    echo "     @home_shared    (新建)"
    echo "  3. 创建测试用户 $TEST_USER"
    echo "  4. 配置共享目录 /home/shared"
    echo "  5. 更新 /etc/fstab"
    if [ "$INSTALL_NIRI" = "yes" ]; then
        echo "  6. 安装 niri"
    fi
    echo ""
    echo "========================================"
    echo ""
    
    log_warn "此操作将修改系统配置，请确保已备份重要数据"
    read -p "确认继续？ [y/N] " -n 1 -r
    echo
    if [[ ! $REPLY =~ ^[Yy]$ ]]; then
        log_info "已取消操作"
        exit 0
    fi
    
    echo "" | tee -a "$LOG_FILE"
    log_info "开始转换..." | tee -a "$LOG_FILE"
    echo "" | tee -a "$LOG_FILE"
}

##############################################################################
# 3. 创建快照
##############################################################################

create_snapshot() {
    log_step "创建 Snapper 快照..." | tee -a "$LOG_FILE"
    
    if ! command -v snapper &> /dev/null; then
        log_warn "跳过快照创建（Snapper 未安装）" | tee -a "$LOG_FILE"
        return
    fi
    
    # 检查 Snapper 配置是否存在
    if ! snapper list-configs | grep -q "root"; then
        log_warn "Snapper 未配置，跳过快照创建" | tee -a "$LOG_FILE"
        log_info "建议先初始化 Snapper: sudo /usr/local/bin/init-snapper.sh" | tee -a "$LOG_FILE"
        return
    fi
    
    SNAPSHOT_ID=$(snapper create --description "转换双用户方案前" --print-number)
    
    if [ -n "$SNAPSHOT_ID" ]; then
        log_success "快照创建成功，ID: $SNAPSHOT_ID" | tee -a "$LOG_FILE"
        log_info "如需回滚: sudo snapper undochange ${SNAPSHOT_ID}..0" | tee -a "$LOG_FILE"
    else
        log_warn "快照创建失败，但将继续执行" | tee -a "$LOG_FILE"
    fi
    
    echo "" | tee -a "$LOG_FILE"
}

##############################################################################
# 4. 重构子卷
##############################################################################

restructure_subvolumes() {
    log_step "重构 Btrfs 子卷..." | tee -a "$LOG_FILE"
    
    # 创建临时挂载点
    TEMP_MNT=$(mktemp -d)
    log_info "临时挂载点: $TEMP_MNT" | tee -a "$LOG_FILE"
    
    # 挂载 Btrfs 根文件系统
    log_info "挂载 Btrfs 根文件系统..." | tee -a "$LOG_FILE"
    mount "$BTRFS_ROOT" "$TEMP_MNT"
    
    # 验证 @home 子卷存在
    if [ ! -d "$TEMP_MNT/@home" ]; then
        log_error "@home 子卷不存在" | tee -a "$LOG_FILE"
        umount "$TEMP_MNT"
        rmdir "$TEMP_MNT"
        exit 1
    fi
    
    # 列出当前子卷
    log_info "当前子卷列表:" | tee -a "$LOG_FILE"
    btrfs subvolume list "$TEMP_MNT" | tee -a "$LOG_FILE"
    echo "" | tee -a "$LOG_FILE"
    
    # 重命名 @home → @home_<主用户>
    log_info "重命名: @home → @home_$MAIN_USER" | tee -a "$LOG_FILE"
    mv "$TEMP_MNT/@home" "$TEMP_MNT/@home_$MAIN_USER"
    log_success "@home 已重命名为 @home_$MAIN_USER" | tee -a "$LOG_FILE"
    
    # 创建 @home_test
    log_info "创建子卷: @home_$TEST_USER" | tee -a "$LOG_FILE"
    btrfs subvolume create "$TEMP_MNT/@home_$TEST_USER"
    log_success "@home_$TEST_USER 创建完成" | tee -a "$LOG_FILE"
    
    # 创建 @home_shared
    log_info "创建子卷: @home_shared" | tee -a "$LOG_FILE"
    btrfs subvolume create "$TEMP_MNT/@home_shared"
    log_success "@home_shared 创建完成" | tee -a "$LOG_FILE"
    
    # 列出新的子卷结构
    echo "" | tee -a "$LOG_FILE"
    log_info "新的子卷列表:" | tee -a "$LOG_FILE"
    btrfs subvolume list "$TEMP_MNT" | tee -a "$LOG_FILE"
    
    # 卸载临时挂载
    log_info "卸载临时挂载点..." | tee -a "$LOG_FILE"
    umount "$TEMP_MNT"
    rmdir "$TEMP_MNT"
    
    log_success "子卷重构完成" | tee -a "$LOG_FILE"
    echo "" | tee -a "$LOG_FILE"
}

##############################################################################
# 5. 更新 fstab
##############################################################################

update_fstab() {
    log_step "更新 /etc/fstab..." | tee -a "$LOG_FILE"
    
    # 备份原 fstab
    cp /etc/fstab /etc/fstab.bak.$(date +%Y%m%d_%H%M%S)
    log_info "已备份 fstab" | tee -a "$LOG_FILE"
    
    # 获取 Btrfs UUID
    BTRFS_UUID=$(blkid -s UUID -o value "$BTRFS_ROOT")
    log_info "Btrfs UUID: $BTRFS_UUID" | tee -a "$LOG_FILE"
    
    # 获取当前挂载选项
    MOUNT_OPTS=$(findmnt -n -o OPTIONS / | grep -oP '(?<=subvol=@,).*' || echo "compress=zstd:3,ssd,noatime,space_cache=v2")
    log_info "挂载选项: $MOUNT_OPTS" | tee -a "$LOG_FILE"
    
    # 修改 @home 的挂载点
    log_info "更新主用户家目录挂载点..." | tee -a "$LOG_FILE"
    sed -i "s|subvol=@home,|subvol=@home_$MAIN_USER,|g" /etc/fstab
    
    # 添加测试用户家目录挂载点
    log_info "添加测试用户家目录挂载点..." | tee -a "$LOG_FILE"
    if ! grep -q "/home/$TEST_USER" /etc/fstab; then
        echo "UUID=$BTRFS_UUID  /home/$TEST_USER  btrfs  subvol=@home_$TEST_USER,$MOUNT_OPTS  0 0" >> /etc/fstab
    fi
    
    # 添加共享目录挂载点
    log_info "添加共享目录挂载点..." | tee -a "$LOG_FILE"
    if ! grep -q "/home/shared" /etc/fstab; then
        echo "UUID=$BTRFS_UUID  /home/shared  btrfs  subvol=@home_shared,$MOUNT_OPTS  0 0" >> /etc/fstab
    fi
    
    # 显示新的 fstab
    echo "" | tee -a "$LOG_FILE"
    log_info "新的 fstab 内容:" | tee -a "$LOG_FILE"
    echo "---" | tee -a "$LOG_FILE"
    grep -E "(/home|btrfs)" /etc/fstab | tee -a "$LOG_FILE"
    echo "---" | tee -a "$LOG_FILE"
    
    log_success "fstab 更新完成" | tee -a "$LOG_FILE"
    echo "" | tee -a "$LOG_FILE"
}

##############################################################################
# 6. 创建测试用户
##############################################################################

create_test_user() {
    log_step "创建测试用户和配置权限..." | tee -a "$LOG_FILE"
    
    # 创建 shared 组
    if ! getent group shared &>/dev/null; then
        log_info "创建 shared 组..." | tee -a "$LOG_FILE"
        groupadd shared
        log_success "shared 组已创建" | tee -a "$LOG_FILE"
    else
        log_info "shared 组已存在" | tee -a "$LOG_FILE"
    fi
    
    # 将主用户加入 shared 组
    log_info "将 $MAIN_USER 加入 shared 组..." | tee -a "$LOG_FILE"
    usermod -aG shared "$MAIN_USER"
    log_success "$MAIN_USER 已加入 shared 组" | tee -a "$LOG_FILE"
    
    # 创建测试用户
    log_info "创建测试用户 $TEST_USER..." | tee -a "$LOG_FILE"
    useradd -m -G wheel,shared -c "测试用户 - niri 环境" "$TEST_USER"
    log_success "用户 $TEST_USER 已创建" | tee -a "$LOG_FILE"
    
    # 设置密码
    echo ""
    log_info "请为用户 $TEST_USER 设置密码:"
    while ! passwd "$TEST_USER"; do
        log_warn "密码设置失败，请重试"
    done
    log_success "密码设置完成" | tee -a "$LOG_FILE"
    
    echo "" | tee -a "$LOG_FILE"
}

##############################################################################
# 7. 配置共享目录
##############################################################################

setup_shared_directory() {
    log_step "配置共享目录..." | tee -a "$LOG_FILE"
    
    # 挂载新的子卷（如果还未挂载）
    log_info "挂载新的子卷..." | tee -a "$LOG_FILE"
    mount -a
    
    # 创建共享目录结构
    log_info "创建共享目录结构..." | tee -a "$LOG_FILE"
    mkdir -p /home/shared/{projects,downloads,documents,media}
    
    # 设置所有者和组
    log_info "设置所有者和权限..." | tee -a "$LOG_FILE"
    chown -R root:shared /home/shared
    chmod -R 2775 /home/shared  # SGID 位
    
    # 验证权限
    log_info "共享目录权限:" | tee -a "$LOG_FILE"
    ls -ld /home/shared | tee -a "$LOG_FILE"
    ls -l /home/shared/ | tee -a "$LOG_FILE"
    
    log_success "共享目录配置完成" | tee -a "$LOG_FILE"
    echo "" | tee -a "$LOG_FILE"
}

##############################################################################
# 8. 安装 niri（可选）
##############################################################################

install_niri() {
    if [ "$INSTALL_NIRI" != "yes" ]; then
        log_info "跳过 niri 安装" | tee -a "$LOG_FILE"
        return
    fi
    
    log_step "安装 niri..." | tee -a "$LOG_FILE"
    
    # 检查 yay
    if ! command -v yay &> /dev/null; then
        log_warn "yay 未安装，跳过 niri 安装" | tee -a "$LOG_FILE"
        log_info "请稍后手动安装: yay -S niri" | tee -a "$LOG_FILE"
        return
    fi
    
    # 以测试用户身份安装（AUR 不允许 root 安装）
    log_info "使用 yay 安装 niri（需要一些时间）..." | tee -a "$LOG_FILE"
    
    # 切换到测试用户安装
    sudo -u "$TEST_USER" yay -S --noconfirm niri || {
        log_warn "niri 安装失败" | tee -a "$LOG_FILE"
        log_info "请稍后手动安装: yay -S niri" | tee -a "$LOG_FILE"
        return
    }
    
    log_success "niri 安装完成" | tee -a "$LOG_FILE"
    
    # 创建基础配置目录
    log_info "创建 niri 配置目录..." | tee -a "$LOG_FILE"
    sudo -u "$TEST_USER" mkdir -p "/home/$TEST_USER/.config/niri"
    
    # 如果存在示例配置，复制过去
    if [ -f /usr/share/niri/niri-config.kdl ]; then
        sudo -u "$TEST_USER" cp /usr/share/niri/niri-config.kdl "/home/$TEST_USER/.config/niri/config.kdl"
        log_info "已创建默认配置文件" | tee -a "$LOG_FILE"
    fi
    
    log_success "niri 配置完成" | tee -a "$LOG_FILE"
    echo "" | tee -a "$LOG_FILE"
}

##############################################################################
# 9. 验证和总结
##############################################################################

verify_and_summary() {
    log_step "验证转换结果..." | tee -a "$LOG_FILE"
    
    # 验证子卷
    echo "" | tee -a "$LOG_FILE"
    log_info "子卷列表:" | tee -a "$LOG_FILE"
    btrfs subvolume list / | grep -E "@home" | tee -a "$LOG_FILE"
    
    # 验证挂载
    echo "" | tee -a "$LOG_FILE"
    log_info "挂载点:" | tee -a "$LOG_FILE"
    findmnt -t btrfs | grep home | tee -a "$LOG_FILE"
    
    # 验证用户
    echo "" | tee -a "$LOG_FILE"
    log_info "用户信息:" | tee -a "$LOG_FILE"
    echo "主用户 $MAIN_USER 组: $(groups $MAIN_USER)" | tee -a "$LOG_FILE"
    echo "测试用户 $TEST_USER 组: $(groups $TEST_USER)" | tee -a "$LOG_FILE"
    
    # 验证共享目录
    echo "" | tee -a "$LOG_FILE"
    log_info "共享目录权限:" | tee -a "$LOG_FILE"
    ls -ld /home/shared | tee -a "$LOG_FILE"
    
    echo "" | tee -a "$LOG_FILE"
    log_success "所有验证通过！" | tee -a "$LOG_FILE"
    
    # 显示总结
    echo "" | tee -a "$LOG_FILE"
    echo "========================================" | tee -a "$LOG_FILE"
    echo "  转换完成！" | tee -a "$LOG_FILE"
    echo "========================================" | tee -a "$LOG_FILE"
    echo "" | tee -a "$LOG_FILE"
    echo "系统配置摘要:" | tee -a "$LOG_FILE"
    echo "  主用户:     $MAIN_USER" | tee -a "$LOG_FILE"
    echo "  测试用户:   $TEST_USER" | tee -a "$LOG_FILE"
    echo "  共享目录:   /home/shared/" | tee -a "$LOG_FILE"
    if [ -n "$SNAPSHOT_ID" ]; then
        echo "  快照 ID:    $SNAPSHOT_ID" | tee -a "$LOG_FILE"
    fi
    echo "" | tee -a "$LOG_FILE"
    echo "子卷结构:" | tee -a "$LOG_FILE"
    echo "  @                → /" | tee -a "$LOG_FILE"
    echo "  @home_$MAIN_USER → /home/$MAIN_USER/" | tee -a "$LOG_FILE"
    echo "  @home_$TEST_USER → /home/$TEST_USER/" | tee -a "$LOG_FILE"
    echo "  @home_shared     → /home/shared/" | tee -a "$LOG_FILE"
    echo "" | tee -a "$LOG_FILE"
    echo "下一步:" | tee -a "$LOG_FILE"
    echo "  1. 重启系统使所有更改生效" | tee -a "$LOG_FILE"
    echo "     sudo reboot" | tee -a "$LOG_FILE"
    echo "" | tee -a "$LOG_FILE"
    echo "  2. 重启后，在 SDDM 登录界面:" | tee -a "$LOG_FILE"
    echo "     - 选择用户: $MAIN_USER 或 $TEST_USER" | tee -a "$LOG_FILE"
    echo "     - 选择会话: Plasma (KDE) 或 niri" | tee -a "$LOG_FILE"
    echo "" | tee -a "$LOG_FILE"
    echo "  3. 测试共享目录:" | tee -a "$LOG_FILE"
    echo "     在任一用户下创建文件:" | tee -a "$LOG_FILE"
    echo "     touch /home/shared/projects/test.txt" | tee -a "$LOG_FILE"
    echo "     切换到另一用户后应能访问" | tee -a "$LOG_FILE"
    echo "" | tee -a "$LOG_FILE"
    
    if [ -n "$SNAPSHOT_ID" ]; then
        echo "如果遇到问题，可以回滚:" | tee -a "$LOG_FILE"
        echo "  sudo snapper undochange ${SNAPSHOT_ID}..0" | tee -a "$LOG_FILE"
        echo "  sudo reboot" | tee -a "$LOG_FILE"
        echo "" | tee -a "$LOG_FILE"
    fi
    
    echo "日志文件: $LOG_FILE" | tee -a "$LOG_FILE"
    echo "" | tee -a "$LOG_FILE"
    echo "========================================" | tee -a "$LOG_FILE"
    echo "" | tee -a "$LOG_FILE"
    
    log_success "转换成功完成！请重启系统。" | tee -a "$LOG_FILE"
}

##############################################################################
# 主流程
##############################################################################

main() {
    # 初始化日志
    echo "===========================================" > "$LOG_FILE"
    echo "双用户方案转换日志" >> "$LOG_FILE"
    echo "时间: $(date)" >> "$LOG_FILE"
    echo "===========================================" >> "$LOG_FILE"
    echo "" >> "$LOG_FILE"
    
    # 执行转换步骤
    check_system
    collect_info
    create_snapshot
    restructure_subvolumes
    update_fstab
    create_test_user
    setup_shared_directory
    install_niri
    verify_and_summary
}

# 执行主流程
main "$@"

